import math
import sys
from polyevaluator import PolyEvaluator 
from gnuplot import GnuPlot
from trainer import Trainer

def main():
  if len(sys.argv) < 5:
    print ('python train_and_plot.py <train_file> '
        '<output-file> <true-file> <regularization>')
    return
  input_file = sys.argv[1]
  output_file = sys.argv[2]
  true_file = sys.argv[3]
  regularization = int(sys.argv[4])

  # Train and get weights
  trainer = Trainer()
  weights = trainer.train([input_file], regularization)

  # Evaluate the learned polynomial
  evaluator = PolyEvaluator(weights)
  learned_poly_filepath = '/tmp/learned.data'
  evaluator.write_data(learned_poly_filepath)

  # Plot data series
  plot_data_series = {}
  plot_data_series['True Data'] = true_file
  plot_data_series['Train Data'] = input_file
  plot_data_series['Learned Polynomial'] = learned_poly_filepath

  gplot = GnuPlot()
  title = 'Lambda = exp(%f)' % regularization
  gplot.plot(plot_data_series, output_file, title)

main()
